"""
向量入库脚本（阶段1）：
- 生成或读取商品数据
- 调用嵌入模型生成向量并做L2归一化
- 创建/修复 Milvus 集合与 COSINE 索引
- 插入数据并 load
运行： python vector_ingest.py
"""
import json
import random
from typing import List, Dict
from pymilvus import utility, FieldSchema, CollectionSchema, DataType, Collection
import numpy as np

from vector_common import (
    OLLAMA_HOST, OLLAMA_PORT, EMBED_MODEL, EMBED_DIM,
    MILVUS_HOST, MILVUS_PORT, COLLECTION_NAME,
    check_tcp, l2_normalize, connect_milvus, ollama_embed
)


def generate_mock_products(n: int = 40) -> List[Dict]:
    categories = ["手机", "电脑", "家电", "服饰", "美妆", "图书", "家居", "运动", "玩具"]
    brands = ["ABrand", "BBrand", "CBrand", "DBrand"]
    products = []
    for i in range(n):
        cat = random.choice(categories)
        brand = random.choice(brands)
        name = f"{brand}{cat}{i+1}号"
        desc = f"来自{brand}的{cat}，高性价比且口碑良好。"
        spec = f"{cat} - {brand} - 型号{i+1} - 颜色随机"
        text = f"名称:{name} 描述:{desc} 规格:{spec} 类别:{cat} 品牌:{brand}"
        products.append({
            "pid": i + 1,
            "name": name,
            "category": cat,
            "brand": brand,
            "text": text
        })
    return products


def ensure_collection(dim: int = EMBED_DIM) -> Collection:
    if utility.has_collection(COLLECTION_NAME):
        coll = Collection(COLLECTION_NAME)
        # 修复索引为 COSINE
        try:
            if coll.indexes:
                idx_params = coll.indexes[0].params or {}
                metric = (idx_params.get("metric_type") or idx_params.get("metric") or "").upper()
                if metric and metric != "COSINE":
                    try:
                        coll.release()
                    except Exception:
                        pass
                    coll.drop_index()
                    index_params = {
                        "index_type": "IVF_FLAT",
                        "metric_type": "COSINE",
                        "params": {"nlist": 1024}
                    }
                    coll.create_index(field_name="embedding", index_params=index_params)
                    coll.load()
            return coll
        except Exception:
            return coll

    fields = [
        FieldSchema(name="pid", dtype=DataType.INT64, is_primary=True, auto_id=False),
        FieldSchema(name="name", dtype=DataType.VARCHAR, max_length=256),
        FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
        FieldSchema(name="brand", dtype=DataType.VARCHAR, max_length=64),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=2048),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim)
    ]
    schema = CollectionSchema(fields=fields, description="Products embeddings")
    coll = Collection(name=COLLECTION_NAME, schema=schema)
    index_params = {"index_type": "IVF_FLAT", "metric_type": "COSINE", "params": {"nlist": 1024}}
    coll.create_index(field_name="embedding", index_params=index_params)
    return coll


def ingest(products: List[Dict], batch: int = 16):
    coll = ensure_collection(EMBED_DIM)

    ids, names, cats, brands, texts = [], [], [], [], []
    for p in products:
        ids.append(p["pid"])
        names.append(p["name"])
        cats.append(p["category"])
        brands.append(p["brand"])
        texts.append(p["text"])

    embeddings: List[List[float]] = []
    for i in range(0, len(texts), batch):
        chunk = texts[i:i+batch]
        vecs = ollama_embed(chunk)
        for v in vecs:
            embeddings.append(l2_normalize(v))

    coll.insert([ids, names, cats, brands, texts, embeddings])
    coll.flush()
    coll.load()


def main():
    assert check_tcp(OLLAMA_HOST, OLLAMA_PORT), f"Ollama not reachable at {OLLAMA_HOST}:{OLLAMA_PORT}"
    assert check_tcp(MILVUS_HOST, int(MILVUS_PORT)), f"Milvus not reachable at {MILVUS_HOST}:{MILVUS_PORT}"
    connect_milvus()

    products = generate_mock_products(40)
    ingest(products)
    print("入库完成：", COLLECTION_NAME)


if __name__ == "__main__":
    main()
